// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/parallel_for.h>
#include "melon/executor.h"
#include "melon/executors/cpu_thread_pool_executor.h"
#include "melon/executors/inline_executor.h"
#include "gmock/gmock.h"
#include <gtest/gtest.h>
#include <pollux/common/base/pollux_exception.h>

using namespace ::testing;
using namespace ::kumo::pollux::dwio::common;

namespace {

class CountingExecutor : public melon::Executor {
 public:
  explicit CountingExecutor(melon::Executor& executor)
      : executor_(executor), count_(0) {}

  void add(melon::Func f) override {
    executor_.add(std::move(f));
    ++count_;
  }

  size_t getCount() const {
    return count_;
  }

 private:
  melon::Executor& executor_;
  size_t count_;
};

void testParallelFor(
    melon::Executor* executor,
    size_t from,
    size_t to,
    size_t parallelismFactor) {
  std::optional<CountingExecutor> countedExecutor;
  std::ostringstream oss;
  oss << "ParallelFor(executor: " << executor << ", from: " << from
      << ", to: " << to << ", parallelismFactor: " << parallelismFactor << ")";
  SCOPED_TRACE(oss.str());
  if (executor) {
    countedExecutor.emplace(*executor);
    executor = &countedExecutor.value();
  }

  std::unordered_map<size_t, std::atomic<size_t>> indexInvoked;
  for (size_t i = from; i < to; ++i) {
    indexInvoked[i] = 0UL;
  }

  ParallelFor(executor, from, to, parallelismFactor)
      .execute([&indexInvoked](size_t i) {
        auto it = indexInvoked.find(i);
        ASSERT_NE(it, indexInvoked.end());
        ++it->second;
      });

  // Parallel For should have thrown otherwise
  ASSERT_LE(from, to);

  // The method was called for each index just once, and didn't call out of
  // bounds indices.
  EXPECT_EQ(indexInvoked.size(), (to - from));
  for (auto& [i, count] : indexInvoked) {
    if (i < from || i >= to) {
      EXPECT_EQ(indexInvoked[i], 0);
    } else {
      EXPECT_EQ(indexInvoked[i], 1);
    }
  }

  if (countedExecutor) {
    const auto extraThreadsUsed = countedExecutor->getCount();
    const auto numTasks = to - from;
    const auto expectedExtraThreads = std::min(
        parallelismFactor > 0 ? parallelismFactor - 1 : 0,
        numTasks > 0 ? numTasks - 1 : 0);
    EXPECT_EQ(extraThreadsUsed, expectedExtraThreads);
  }
}

} // namespace

TEST(ParallelForTest, E2E) {
  auto inlineExecutor = melon::InlineExecutor::instance();
  for (size_t parallelism = 0; parallelism < 25; ++parallelism) {
    for (size_t begin = 0; begin < 25; ++begin) {
      for (size_t end = 0; end < 25; ++end) {
        if (begin <= end) {
          testParallelFor(&inlineExecutor, begin, end, parallelism);
        } else {
          EXPECT_THROW(
              testParallelFor(&inlineExecutor, begin, end, parallelism),
              kumo::pollux::PolluxRuntimeError);
        }
      }
    }
  }
}

TEST(ParallelForTest, E2EParallel) {
  for (size_t parallelism = 1; parallelism < 2; ++parallelism) {
    melon::CPUThreadPoolExecutor executor(parallelism);
    for (size_t begin = 0; begin < 25; ++begin) {
      for (size_t end = 0; end < 25; ++end) {
        if (begin <= end) {
          testParallelFor(&executor, begin, end, parallelism);
        } else {
          EXPECT_THROW(
              testParallelFor(&executor, begin, end, parallelism),
              kumo::pollux::PolluxRuntimeError);
        }
      }
    }
  }
}

TEST(ParallelForTest, CanOwnExecutor) {
  auto executor = std::make_shared<melon::CPUThreadPoolExecutor>(2);
  const size_t indexInvokedSize = 100;
    std::unordered_map<size_t, std::atomic<size_t>> indexInvoked;
  indexInvoked.reserve(indexInvokedSize);
  for (size_t i = 0; i < indexInvokedSize; ++i) {
    indexInvoked[i] = 0UL;
  }

  ParallelFor pf(executor, 0, indexInvokedSize, 9);
  pf.execute([&indexInvoked](size_t i) { ++indexInvoked[i]; });

  EXPECT_EQ(indexInvoked.size(), indexInvokedSize);
  for (size_t i = 0; i < indexInvokedSize; ++i) {
    EXPECT_EQ(indexInvoked[i], 1);
  }
}
